package com.lz.service.config;

import com.lz.common.baseBean.R;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import javax.servlet.http.HttpServletRequest;
import java.util.Date;

//@RestControllerAdvice
//public class TraceIdResponseAdvice implements ResponseBodyAdvice<Object> {
//
//    private final HttpServletRequest request;
//
//    public TraceIdResponseAdvice(HttpServletRequest request) {
//        this.request = request;
//    }
//
//    @Override
//    public boolean supports(MethodParameter returnType, Class converterType) {
//        return true;
//    }
//
//    @Override
//    public Object beforeBodyWrite(Object body, MethodParameter returnType, MediaType selectedContentType, Class selectedConverterType, ServerHttpRequest request, ServerHttpResponse response) {
//        if (body instanceof Map) {
//            String traceId = (String) this.request.getAttribute("traceId");
//            ((Map<String, Object>) body).put("traceId", traceId);
//        } else if (body instanceof R) {
//            String traceId = (String) this.request.getAttribute("traceId");
//            ((R<?>) body).put( traceId);
//        }
//        return body;
//    }
//}

@RestControllerAdvice
public class TraceIdResponseAdvice implements ResponseBodyAdvice<R<?>> {

    @Autowired
    private HttpServletRequest request;

    @Override
    public boolean supports(MethodParameter returnType, Class converterType) {
        return true;
    }


    @Override
    public R<?> beforeBodyWrite(R<?> r, MethodParameter methodParameter, MediaType mediaType, Class<? extends HttpMessageConverter<?>> aClass, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {
        String traceId = (String) request.getAttribute("traceId");
        if (traceId != null && r != null) {
            r.addTraceIdAndTime(traceId, new Date());
        }

        return r;
    }
}